"""证明材料生成相关节点"""

import asyncio
import uuid
from pocketflow import AsyncParallelBatchNode
from ..utils.call_llm import call_llm_async
from ..utils.proof_mapper_utils import create_materials_prompt, parse_materials_response

class ProofMapperNode(AsyncParallelBatchNode):
    """并行为所有评估细项生成证明材料清单"""
    
    def __init__(self, max_concurrency=15, max_retries=3, **kwargs):
        super().__init__(**kwargs)
        self.semaphore = asyncio.Semaphore(max_concurrency)  # 15线程并发控制
        self.max_retries = max_retries
    
    async def prep_async(self, shared):
        checkpoints = shared.get("checkpoints", [])
        if not checkpoints:
            print("警告: 没有评估项可处理")
            return []
        
        print(f"准备并行生成 {len(checkpoints)} 个评估项的证明材料清单")
        return checkpoints
    
    async def exec_async(self, checkpoint):
        """为单个评估细项生成证明材料清单"""
        async with self.semaphore:  # 控制并发数
            # 创建提示词并调用LLM
            prompt = create_materials_prompt(checkpoint)
            
            # 添加重试机制
            for retry in range(self.max_retries):
                try:
                    response = await call_llm_async(prompt)
                    return parse_materials_response(response, checkpoint)
                    
                except Exception as e:
                    if retry < self.max_retries - 1:
                        print(f"生成证明材料时出错 (重试 {retry+1}/{self.max_retries}): {str(e)}")
                        await asyncio.sleep(1)  # 等待1秒后重试
                    else:
                        print(f"生成证明材料失败，达到最大重试次数: {str(e)}")
                        return {
                            "id": str(uuid.uuid4()),
                            "module": checkpoint["module"],
                            "category": checkpoint["category"],
                            "checkpoint": checkpoint["checkpoint"],
                            "materials": ["处理失败，请手动添加证明材料"],
                            "notes": f"处理错误: {str(e)}"
                        }
    
    async def post_async(self, shared, prep_res, exec_res_list):
        # 检查是否所有评估项都被处理
        processed_items = len(exec_res_list)
        total_items = len(prep_res)
        if processed_items != total_items:
            print(f"警告: 只处理了 {processed_items}/{total_items} 个评估项")
            
        shared["materials"] = exec_res_list
        print(f"并行生成完成，共生成了 {len(exec_res_list)} 个证明材料清单")
        return "default"
